from torchvision import transforms
from torch.utils import data
from PIL import Image

def load_batch_test(config, filename):
    img = Image.open(filename).convert('RGB')
    im_size = img.size
    img = transforms.Resize((config.input_size, config.input_size))(img)
    img = transforms.ToTensor()(img)
    img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img)
    return img, im_size


def get_loader(filename, config, mode = 'train'):
    dataset = ImageDataTest(filename, config)
    data_loader = data.DataLoader(dataset = dataset, batch_size = 1, shuffle = False, num_workers = 0)
    return data_loader


class ImageDataTest(data.Dataset):
    def __init__(self, filename, config):
        self.filename = filename
        self.image_list = [filename]
        self.config = config
        self.image_num = len(self.image_list)

    def __getitem__(self, item):
        image, im_size = load_batch_test(self.config, self.filename)
        return {'image': image, 'name': self.image_list[item % self.image_num], 'size': im_size}

    def __len__(self):
        return self.image_num